import os
import cv2
import sys
import utils
import random
import numpy as np
from tqdm import tqdm


def TransparentEdge(img):
    # minLineLength = img.shape[0] / 32  # height/32
    # maxLineGap = img.shape[0] / 40  # height/40
    TopLeft = [100, 100]
    TopRight = [0, 100]
    BottomLeft = [100, 0]
    BottomRight = [0, 0]
    AngularPoint = [[100, 100], [0, 100], [100, 0], [0, 0]]

    minLineLength = 2
    maxLineGap = 2
    result = img.copy()
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # gray = np.float32(gray)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, 10, minLineLength, maxLineGap)
    # 输入图像必须是 float32 ,最后一个参数在 0.04 到 0.05 之间
    # dst = cv2.cornerHarris(gray, 2, 3, 0.04)
    for line in lines:
        for x1, y1, x2, y2 in line:
            if x1 < TopLeft[0] and y1 < TopLeft[1]:
                TopLeft[0] = x1
                TopLeft[1] = y1
            if x2 < TopLeft[0] and y2 < TopLeft[1]:
                TopLeft[0] = x2
                TopLeft[1] = y2
            if x1 > TopRight[0] and y1 < TopRight[1]:
                TopRight[0] = x1
                TopRight[1] = y1
            if x2 > TopRight[0] and y2 < TopRight[1]:
                TopRight[0] = x2
                TopRight[1] = y2
            if x1 < BottomLeft[0] and y1 > BottomLeft[1]:
                BottomLeft[0] = x1
                BottomLeft[1] = y1
            if x2 < BottomLeft[0] and y2 > BottomLeft[1]:
                BottomLeft[0] = x2
                BottomLeft[1] = y2
            if x1 > BottomRight[0] and y1 > BottomRight[1]:
                BottomRight[0] = x1
                BottomRight[1] = y1
            if x2 > BottomRight[0] and y2 > BottomRight[1]:
                BottomRight[0] = x2
                BottomRight[1] = y2
    cv2.line(result, (TopLeft[0], TopLeft[1]), (TopRight[0], TopRight[1]), (0, 255, 0), 2)
    cv2.line(result, (TopLeft[0], TopLeft[1]), (BottomLeft[0], BottomLeft[1]), (0, 255, 0), 2)
    cv2.line(result, (BottomLeft[0], BottomLeft[1]), (BottomRight[0], BottomRight[1]), (0, 255, 0), 2)
    cv2.line(result, (TopRight[0], TopRight[1]), (BottomRight[0], BottomRight[1]), (0, 255, 0), 2)
    cv2.imshow('result', result)
    cv2.waitKey(0)
    return img


BalanceDir = r"I:\Temp"  # 需要平衡标签的文件夹
BackGroundPicDir = r"I:\BackGroundPic\road"  # 没有标签的干净图片（用于P图的图片）
DontNeedBalance = []  # 不需要平衡的label
BalanceNum = 500  # 0表示平衡到数据量最多的标签
CutRateRange = [1, 1.5]

FileNames = []
BackGroundPicPaths = []
NameNum = 0


print("open background dir...")
for root, dirs, files in os.walk(BackGroundPicDir):
    for file in files:
        if file[-1] == 'g':
            BackGroundPicPaths.append(file)
# print("Count label types...")
LabelDict = utils.CountLabelKind(BalanceDir)
for DontNeedBalanceLabel in DontNeedBalance:
    del LabelDict[DontNeedBalanceLabel]
if BalanceNum == 0:
    BalanceNum = max(LabelDict.values())
print("pictures classified by label...")
for key in LabelDict.keys():
    FilesList = []
    for root, dirs, files in os.walk(BalanceDir):
        for file in tqdm(files):
            if file[-1] == 'l':
                Infos = utils.ReadXml(root + "\\" + file)
                for Info in Infos:
                    if Info[-1] == key and file not in FilesList:
                        FilesList.append(file)
    FileNames.append(FilesList)
print("star to generate picture...")
for i, key in enumerate(LabelDict.keys()):
    while LabelDict[key] < BalanceNum:
        Num = random.randint(0, len(FileNames[i])-1)
        Img = cv2.imread(BalanceDir + "\\" + FileNames[i][Num][:-3] + 'jpg')
        Infos = utils.ReadXml(BalanceDir + "\\" + FileNames[i][Num])

        if len(Infos) == 1:
            for Info in Infos:
                if Info[-1] == key:
                    ImgLabelWithBackground, IndexLabelWithBackground, _ = utils.CutLabelWithBackground(Img, np.array(Infos)[:, :-1].astype(int), CutRateRange, Info[0], Info[1], Info[2], Info[3])
                    # for xmin, ymin, xmax, ymax in IndexLabelWithBackground:
                    #     cv2.rectangle(ImgLabelWithBackground, (xmin, ymin), (xmax, ymax), (0, 0, 255), 1)
                    # cv2.imshow("test", ImgLabelWit)
                    LabelImg = ImgLabelWithBackground
                    # InfoLabelWithBackground = (np.insert(IndexLabelWithBackground, 4, values=np.array(Infos)[:, -1], axis=1)).tolist()[0]

                    # LabelImg = Img[Info[1]:Info[3], Info[0]:Info[2]]
                    # 随机选择背景图
                    Num = random.randint(0, len(BackGroundPicPaths) - 1)
                    BackGroundImg = cv2.imread(BackGroundPicDir + "\\" + BackGroundPicPaths[Num])
                    BackGroundImg = cv2.resize(BackGroundImg, (1920, 1080))
                    LabelImg = utils.BrightnessAverage(BackGroundImg, LabelImg)

                    # LabelImg = TransparentEdge(LabelImg)  # 函数未完成

                    RandW = random.randint(0, BackGroundImg.shape[1] - LabelImg.shape[1] - 2)
                    RandH = random.randint(0, BackGroundImg.shape[0] - LabelImg.shape[0] - 2)
                    BackGroundImg[RandH:RandH + LabelImg.shape[0], RandW:RandW + LabelImg.shape[1]] = LabelImg
                    # cv2.imshow("test", CleanImg)
                    # cv2.waitKey(0)
                    NameNum = NameNum + 1
                    NewInfos = [BalanceDir + "\\Balance%05d.xml" % NameNum]
                    NewInfo = [RandW + IndexLabelWithBackground[0][0], RandH + IndexLabelWithBackground[0][1], RandW + IndexLabelWithBackground[0][2], RandH + IndexLabelWithBackground[0][3], Info[-1]]
                    NewInfos.append(NewInfo)
                    print(NewInfos[0])
                    utils.WriteXml(NewInfos, BackGroundImg.shape[1], BackGroundImg.shape[0])
                    cv2.imwrite(BalanceDir + "\\Balance%05d.jpg" % NameNum, BackGroundImg)
                    LabelDict[key] = LabelDict[key] + 1
                    break


